
# Simulation of atmospheric vertical wind plus settling
#
import numpy as np  # Import numerical python (arrays vectorization etc.)
from scipy.optimize import fsolve
from math import *
import matplotlib.pyplot as plt # Plotting
from math import log as ln
from math import log10 as log
import time

from PyQt5 import QtGui  # Gui input interface
from PyQt5 import QtCore
import pyqtgraph as pg
import re

print('  ')
print('  ')
print ('  +---------------------------------------+')
print ('  |         EVAPORATION WIND MODEL        |')
print ('  |                                       |')
print ('  |     Includes isothermal Parker wind   |')
print ('  |                                       |')
print ('  |           (EdY & HT  June, 2020)      |')
print ('  +---------------------------------------+')
print('')


# FRONT MATTER FOR GUI INPUT
# -----------------------------------------------------------------------
# SCIENTIFIC NOTATION version of QDoubleSpinBox()
# Regular expression to find floats. Match groups are the whole string, the
# whole coefficient, the decimal part of the coefficient, and the exponent
# part.
_float_re = re.compile(r'(([+-]?\d+(\.\d*)?|\.\d+)([eE][+-]?\d+)?)')

def valid_float_string(string):
    match = _float_re.search(string)
    return match.groups()[0] == string if match else False


class FloatValidator(QtGui.QValidator):

    def validate(self, string, position):
        if valid_float_string(string):
            state = QtGui.QValidator.Acceptable
        elif string == "" or string[position-1] in 'e.-+':
            state=QtGui.QValidator.Intermediate
        else:
            state=QtGui.QValidator.Invalid
        return (state,string,position)

    def fixup(self, text):
        match = _float_re.search(text)
        return match.groups()[0] if match else ""


class ScientificDoubleSpinBox(QtGui.QDoubleSpinBox):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.setMinimum(-np.inf)
        self.setMaximum(np.inf)
        self.validator = FloatValidator()
        self.setDecimals(1000)
        self.setValue(np.inf)
        self.setSingleStep(0.1)

    def validate(self, text, position):
        return self.validator.validate(text, position)

    def fixup(self, text):
        return self.validator.fixup(text)

    def valueFromText(self, text):
        return float(text)

    def textFromValue(self, value):
        return format_float(value)

    def stepBy(self, steps):
        text = self.cleanText()
        groups = _float_re.search(text).groups()
        decimal = float(groups[1])
        decimal += steps
        new_string = "{:g}".format(decimal) + (groups[3] if groups[3] else "")
        self.lineEdit().setText(new_string)


def format_float(value):
    """Modified form of the 'g' format specifier."""
    string = "{:g}".format(value).replace("e+", "e")
    string = re.sub("e(-?)0*(\d+)", r"e\1\2", string)
    return string
# END scientific notation version of QDoubleSpinBox()


# BEGIN GUI interface itself
## Always start by initializing Qt (only once per application)
app = QtGui.QApplication([])

## Define a top-level widget to hold everything
w = QtGui.QWidget()

## Create some widgets to be placed inside
label2 = QtGui.QLabel('Enter Temperature')
box2 = QtGui.QDoubleSpinBox()
box2.setMinimum(50.0)
box2.setSingleStep(10.0)
box2.setMaximum(10000.00)
box2.setSuffix("  K")
box2.setValue(1550.0)

label3 = QtGui.QLabel('Keddy')
box3 = QtGui.QDoubleSpinBox()
box3.setSingleStep(10.0)
box3.setMinimum(0.0)
box3.setMaximum(100000000.00)
box3.setSuffix("  m^2/s")
box3.setValue(100.0)

label4 = QtGui.QLabel('Mass of Body')
box4 = ScientificDoubleSpinBox()
box4.setSingleStep(0.1)
box4.setMinimum(0.0001)
box4.setMaximum(1.0e35)
box4.setSuffix(" kg")
box4.setValue(7.35e22)

label15 = QtGui.QLabel('Time Step Size')
box15 = QtGui.QDoubleSpinBox()
box15.setSingleStep(10.0)
box15.setMinimum(0.0)
box15.setMaximum(10000000000.00)
box15.setSuffix(" sec")
box15.setValue(100000.0)

label17 = QtGui.QLabel('Radius Carrier Molec')
box17 = QtGui.QDoubleSpinBox()
box17.setMinimum(0.000)
box17.setMaximum(9000.0)
box17.setSuffix(" pm")
box17.setValue(227.0)

label18 = QtGui.QLabel('Radius Tracer Molec')
box18 = QtGui.QDoubleSpinBox()
box18.setMinimum(0.000)
box18.setMaximum(9000.0)
box18.setSuffix(" pm")
box18.setValue(280.0)

label5 = QtGui.QLabel('Radius of Body')
box5 = QtGui.QDoubleSpinBox()
box5.setMinimum(50.0)
box5.setMaximum(1000000.00)
box5.setSingleStep(10.0)
box5.setSuffix("  km")
box5.setValue(1700.00)

label6 = QtGui.QLabel('Height of Homopause')
box6 = QtGui.QDoubleSpinBox()
box6.setMinimum(0.0)
box6.setMaximum(10000000.00)
box6.setSingleStep(10.0)
box6.setSuffix("  km")
box6.setValue(147.00)

label7 = QtGui.QLabel('Homopause Vap Press')
box7 = ScientificDoubleSpinBox()
box7.setSingleStep(0.5)
box7.setMinimum(1.0e-15)
box7.setMaximum(1.0e00)
box7.setSuffix("  bar")
box7.setValue(1.001e-08)

label8 = QtGui.QLabel('Mass of Carrier Gas')
box8 = QtGui.QDoubleSpinBox()
box8.setMinimum(1.0)
box8.setMaximum(500.00)
box8.setSingleStep(1.0)
box8.setSuffix("  amu")
box8.setValue(23.00)

label9 = QtGui.QLabel('Mass of tracer Major Isotope')
box9 = QtGui.QDoubleSpinBox()
box9.setMinimum(1.0)
box9.setMaximum(500.00)
box9.setSingleStep(1.0)
box9.setSuffix("  amu")
box9.setValue(39.00)

label10 = QtGui.QLabel('Mass of tracer Rare Isotope')
box10 = QtGui.QDoubleSpinBox()
box10.setMinimum(1.0)
box10.setMaximum(500.00)
box10.setSingleStep(1.0)
box10.setSuffix("  amu")
box10.setValue(41.00)

label11 = QtGui.QLabel('Molar tracer/carrier')
box11 = QtGui.QDoubleSpinBox()
box11.setDecimals(4)
box11.setMinimum(0.0001)
box11.setMaximum(100.00)
box11.setSingleStep(0.001)
box11.setValue(0.10)

label12 = QtGui.QLabel('Tracer minor/major')
box12 = QtGui.QDoubleSpinBox()
box12.setMinimum(0.0)
box12.setMaximum(100.00)
box12.setSingleStep(0.0005)
box12.setDecimals(4)
box12.setValue(0.0720)

label13 = QtGui.QLabel('Number of Time Steps')
box13 = QtGui.QSpinBox()
box13.setMinimum(0.0)
box13.setMaximum(100000.00)
box13.setSingleStep(1)
box13.setValue(50)

label14 = QtGui.QLabel('Vel Multiplier at U Bound')
box14 = QtGui.QDoubleSpinBox()
box14.setMinimum(-100.0)
box14.setMaximum(100.0)
box14.setSingleStep(0.01)
box14.setValue(1.0)

label19 = QtGui.QLabel('Body-neighbor distance')
box19 = QtGui.QDoubleSpinBox()
box19.setMinimum(4.9)
box19.setMaximum(100.0)
box19.setSingleStep(1.0)
box19.setValue(100.0)
box19.setSuffix( "% of E-M today")

label20 = QtGui.QLabel('Mass of Neighbor')
box20 = ScientificDoubleSpinBox()
box20.setSingleStep(0.1)
box20.setMinimum(0.0001)
box20.setMaximum(1.0e35)
box20.setSuffix(" kg")
box20.setValue(5.970E24)

b1 = QtGui.QPushButton("Enter")
b1.setStyleSheet('background-color: grey; color: white')

## Create a grid layout to manage the widgets size and position
layout = QtGui.QGridLayout()
w.setLayout(layout)

## Add widgets to the layout in their proper positions
layout.addWidget(label2, 0, 0)   # Column 1
layout.addWidget(box2, 1, 0)   #
layout.addWidget(label3, 2, 0)  #
layout.addWidget(box3,3,0)
layout.addWidget(label4, 4, 0)  #
layout.addWidget(box4, 5, 0)  #
layout.addWidget(label5, 6, 0)  #
layout.addWidget(box5, 7, 0)  #
layout.addWidget(label6, 8, 0)  #
layout.addWidget(box6, 9, 0)  #
layout.addWidget(label7, 10, 0)  #
layout.addWidget(box7, 11, 0)  #

layout.addWidget(label8, 0, 1)  # Column 2
layout.addWidget(box8, 1, 1)  #
layout.addWidget(label9, 2,1)  #
layout.addWidget(box9, 3, 1)  #
layout.addWidget(label10, 4, 1)  #
layout.addWidget(box10, 5, 1)  #
layout.addWidget(label11, 6, 1)  #
layout.addWidget(box11, 7, 1)  #
layout.addWidget(label12, 8, 1)  #
layout.addWidget(box12, 9, 1)  #
layout.addWidget(label14, 10, 1)  #
layout.addWidget(box14, 11, 1)  #

layout.addWidget(label17, 0, 2) #Column 3
layout.addWidget(box17, 1, 2)
layout.addWidget(label18, 2, 2)
layout.addWidget(box18, 3, 2)
layout.addWidget(label15, 4, 2)  #
layout.addWidget(box15, 5, 2)  #
layout.addWidget(label13,6,2) # Number of time steps
layout.addWidget(box13,7,2) # Number of time steps
layout.addWidget(label19,8,2)
layout.addWidget(box19,9,2)
layout.addWidget(label20,10,2)
layout.addWidget(box20,11,2)

layout.addWidget(b1, 1,3) # Enter button third column

b1.clicked.connect(QtCore.QCoreApplication.instance().quit)

## Display the widget as a new window
w.show()

## Start the Qt event loop
app.exec_()

print(' ')
print(' ')

TK = box2.value()
print("Entered Temperature = ",TK,' K')

keddy=box3.value()
print("Entered Keddy = ", keddy,' m^2/s')

Mp = box4.value()
print("Entered Mass = ", Mp,' kg')

R_s = box5.value()
print("Entered Radius = ", R_s,' km')

Z_homo=box6.value()
print("Entered Z homopause = ", Z_homo,' km')

PEQ=box7.value()
print("Entered Eq vap pressure = ","{:.5e}".format(PEQ)," bar")

m_main=box8.value()
print("Entered mass of carrier = ",m_main,' amu')

m_i=box9.value()
print("Entered mass of tracer i = ",m_i,' amu')

m_j=box10.value()
print("Entered mass tracer j = ",m_j,' amu')

mix_i=box11.value()
print("Entered mixing ratio of i = ",mix_i)

mix_j=box12.value()
print("Entered mixing ratio of i = ",mix_j)

nlimit=box13.value()
print("Entered number of time steps = ",nlimit)

vel_factor=box14.value()
print("Entered multiple of velocity for upper boundary = ", vel_factor)

dt=box15.value()
print('Entered time step = ',dt,' seconds')

r_main=box17.value()
r_main=r_main*1.0e-12
print('Entered radius carrier molecule = ',r_main,' meters')

r_i=box18.value()
r_i=r_i*1.0e-12
print('Entered radius tracer molecule = ', r_i,' meters')

EM_relative_D=box19.value()
print('Entered Earth-Moon distance = ', EM_relative_D,' %')

ME=box20.value()
print('Entered nearest neighbor mass = ',ME,' km')



# TEST OF WHAT HAPPENS WHEN THERE ARE FEW COLLISIONS
#r_main=r_main*1.0e-15
#r_i=r_i*1.0e-15


# END GUI INPUT VALUES
# -----------------------------------------------------------------------

# CONSTANTS IN SI UNITS
G=6.67e-11 # m^3 kg^-1 s^-2
nucleon_mass = 1.6726e-27  # kg
kb = 1.38065e-23 # J/K
pi = 3.14159265
Av_number=6.02214e23 # particles per mole

# INPUTS SECTION
# --------------------------------------------------------
print(" ")
print("ENTERED PARAMETERS...")
print(" Atm temperature = ", TK, "K")
print(" Equilibrium vapor pressure = ", PEQ," bar")
#Pbackground=float(input("Enter background pressure (bar) at surface: "))
#Pbackground = 1.0e-16  # bar 1.0e-15

Pbackground=PEQ


m_main_string=str(int(m_main))
mi_string=str(int(m_i))
mj_string=str(int(m_j))
# --------------------------------------------------------
# END BASIC INPUTS

# MODIFY UNITS OF INPUTS AND ECHO TO USER
mix=1.0  #legacy used to correct P for main gas to mole fraction

print(" background species mass = ", m_main," amu")
print(" tracer species mass i = ", m_i," amu")
print(" tracer species mass j = ", m_j," amu")
print(" background species radius = ",r_main," meters" )
print(' ')
print(" tracer mixing ratio for i = ",mix_i)
print(" tracer mixing ratio for j = ",mix_j)

# Some required conversions
m_main=m_main*nucleon_mass
m_i=m_i*nucleon_mass
m_j=m_j*nucleon_mass
ndens_initial = Pbackground*1.0e5/(kb*TK)  # convert bar to number density through Pascal


R_s = R_s*1000.0
print(" Radius of body = ", "{:.5e}".format(R_s)," meters")
print(" Mass of body = ", Mp," kg")
print(" Elevation of homopause = ", Z_homo," km")
Z_homo=Z_homo*1000.0 # convert to meters

# SOUND SPEED
gamma_gas=1.666
#gamma_gas=1.00
# SOUND SPEED
Cs=np.sqrt(gamma_gas*kb*TK/m_main)
Cs_i=np.sqrt(gamma_gas*kb*TK/m_i)
Cs_j=np.sqrt(gamma_gas*kb*TK/m_j)
#Cs_parameter=G*Mp/2.0/Cs**2.0/1000.0 #km
print('')
print("ESCAPE FACTORS")
print(" Cs = ", Cs)


# CRITICAL RADIUS INCLUDING THE EFFECT OF PROXIMAL PLANET
#ME=5.97E24 #mass of Earth
rs=G*Mp/(2.0*Cs**2.0)  # critical radius for no Earth
Cs_parameter=rs/1000.0 # Cs in km units
EM_D_today=384400.0 #km, present EM distance
EM_D=(EM_relative_D/100.0)*EM_D_today
alpha=(Mp/3.0/ME)**(1.0/3.0)
alpha_rev=1.0/alpha
beta=-1.0
Q=EM_D*alpha/Cs_parameter

A5=alpha_rev**3.0/3.0
A4=alpha_rev**4.0*beta
A3=alpha_rev**5.0+Q*alpha_rev**3.0
A2=beta*2.0*Q*alpha_rev**4-alpha_rev**3.0
A1=Q*alpha_rev**5.0-2.0*beta*alpha_rev**4.0
A0=-alpha_rev**5.0

def criticalRadius(zz):
   x = zz[0]
   F = np.empty((1))
   F[0] = A5*x**5.0+A4*x**4.0+A3*x**3.0+A2*x**2.0+A1*x+A0
   return F

zGuess = np.array([1])
zz = fsolve(criticalRadius,zGuess)
r_c_R_h_ratio=zz[0]
r_c_r_s_ratio=Q*r_c_R_h_ratio
rc=r_c_r_s_ratio*rs/1000.0   #new critical radius in km
heightrc=rc-R_s/1000.0
print(' Transonic radius for isolated body =',"{:.2f}".format(rs/1000),'km;')
print(' Transonic radius with distance to neighbor of ',"{:.1f}".format(EM_D),'km is',"{:.2f}".format(rc),'km')
print('   located ',"{:.2f}".format(heightrc),'km above the surface.')
rc = rc*1000.0 #convert to meters
rc_ratio=rc/R_s
# Double checking values for use below.
print(" rc/r_planet = ",rc_ratio)
print(" rc = ", rc," meters")

if heightrc < 0.0:
    print(' *CRITICAL RADIUS BELOW SURFACE, SETTING rc TO HOMOPAUSE + 1km')
    rc=R_s+Z_homo+1000.0
    heightrc=(rc-R_s)/1000.0
    rc_ratio=rc/R_s
    time.sleep(8)
    print(" New rc height = ", heightrc," km")
    time.sleep(5)
# END CALCULATION OF CRITICAL RADIUS


# Escape parameter lambda
lambda_e = (G*Mp*m_main/R_s)/(kb*TK)
print(" Lambda = ",lambda_e)
print('')
lambda_string=str(lambda_e)

# RADIAL SPATIAL DIMENSIONS AND INCREMENTS
rmin=R_s+Z_homo
nr=4000.0  # NUMBER OF RADIAL POSITIONS, 4000.0
n=int(nr)
#rmax = rmin+nr*dr
rmax=rc+0.1*rc
maxht=(rmax-R_s)/1000.0  #for reporting in km
dr=(rmax-rmin)/nr
print(" ")
print("NUMERICAL SETUP...")
print(" Miniumum radius =","{:.5e}".format(rmin)," meters")
print(" Maximum radius = ", "{:.5e}".format(rmax)," meters")
print(" Maximum altitutde = ", maxht," km")
print(" dr = ", dr)
print(" nr = ", int(nr))

# RADIUS ARRAY
r=np.linspace(rmin, rmax, int(nr)+1)
n_array=len(r)
print(" Length of r( ) array = ", n_array," and should be nr+1 to account for 0 index")
#Check dr is correct
dr = (r[1]-r[0])
print(" r[0] = ",r[0])
print(" r[n] = ",r[int(nr)])
print(" Confirm that dr = ", dr)

# HEIGHTS
z_atm=(r-R_s)
height=z_atm/1000.0

# VELOCITY ARRAY
u_h=np.linspace(rmin,rmax,n+1)
u_h_i=np.linspace(rmin,rmax,n+1)
u_h_j=np.linspace(rmin,rmax,n+1)

# NEWTONS METHOD FOR HYDRODYNAMIC FLOW VELOCITIES for carrier
NREPORT=100
for i in range(0,n):
    precision = 1.0e-7
    if r[i] < rc:
        utest=1.0e-6
    else:
        utest = Cs+5.0
    k=0
    s=0
    for k in range (1,100):
        f=(utest**2.0/Cs**2.0)-ln(utest**2.0/Cs**2.0)+3.0-4.0*ln(r[i]/rc)-4.0*rc/r[i]
        fprime=(2.0*utest/Cs**2)-2.0/utest
        if abs(f) < precision:
            u_h[i]=utest
            s=k
            break
#            if i%NREPORT == 0:
#                print("...successfull solution for radius ",i)
#                print(".....at step ",k)
#                print("U = ",u[i])
        else:
            utest=utest-f/fprime
        k=k+1
#    print('Steps for solution = ',s,' at radius ',i)
#    print(' fprime = ',fprime)

# NEWTONS METHOD FOR HYDRODYNAMIC FLOW VELOCITIES for species i
NREPORT=100
for i in range(0,n):
    precision = 1.0e-7
    if r[i] < rc:
        utest=1.0e-6
    else:
        utest = Cs+5.0
    k=0
    s=0
    for k in range (1,100):
        f=(utest**2.0/Cs_i**2.0)-ln(utest**2.0/Cs_i**2.0)+3.0-4.0*ln(r[i]/rc)-4.0*rc/r[i]
        fprime=(2.0*utest/Cs_i**2)-2.0/utest
        if abs(f) < precision:
            u_h_i[i]=utest
            s=k
            break
#            if i%NREPORT == 0:
#                print("...successfull solution for radius ",i)
#                print(".....at step ",k)
#                print("U = ",u[i])
        else:
            utest=utest-f/fprime
        k=k+1
#    print('Steps for solution = ',s,' at radius ',i)
#    print(' fprime = ',fprime)

# NEWTONS METHOD FOR HYDRODYNAMIC FLOW VELOCITIES for species j
NREPORT=100
for i in range(0,n):
    precision = 1.0e-7
    if r[i] < rc:
        utest=1.0e-6
    else:
        utest = Cs+5.0
    k=0
    s=0
    for k in range (1,100):
        f=(utest**2.0/Cs_j**2.0)-ln(utest**2.0/Cs_j**2.0)+3.0-4.0*ln(r[i]/rc)-4.0*rc/r[i]
        fprime=(2.0*utest/Cs_j**2)-2.0/utest
        if abs(f) < precision:
            u_h_j[i]=utest
            s=k
            break
#            if i%NREPORT == 0:
#                print("...successfull solution for radius ",i)
#                print(".....at step ",k)
#                print("U = ",u[i])
        else:
            utest=utest-f/fprime
        k=k+1
#    print('Steps for solution = ',s,' at radius ',i)
#    print(' fprime = ',fprime)


# AREAS AND VOLUMES
area=np.zeros(n)
area=4.0*pi*r**2.0
volume=np.zeros(n)
volume[0]=(4.0/3.0)*pi*(r[0]+dr)**3.0-(4.0/3.0)*pi*(r[0])**3.0
for i in range(1,n-1):
    volume=((4.0/3.0)*pi*(r+dr)**3.0-(4.0/3.0)*pi*(r-dr)**3.0 )/2.0
volume[n]=(4.0/3.0)*pi*(r[n])**3.0-(4.0/3.0)*pi*(r[n]-dr)**3.0


# BINARY DIFFUSION COEFFICIENTS
xsection=pi*((r_i+r_main)/1.0)**2.0
print('')
print("DIFFUSION PARAMETERS...")
print(" Collision cross section = ", "{:.5e}".format(xsection)," meters")
red_mass_i=m_i*m_main/(m_i+m_main)
red_mass_j=m_j*m_main/(m_j+m_main)
v_avg_i=sqrt(kb*TK/red_mass_i)
v_avg_j=sqrt(kb*TK/red_mass_j)
print(" Avg thermal velocity i vs buffer gas = ", v_avg_i," m/s")
print(" Avg thermal velocity j vs buffer gas = ", v_avg_j," m/s")
b_binary_i=3.0*sqrt(2.0*pi)*v_avg_i/(64.0*xsection)
b_binary_j=3.0*sqrt(2.0*pi)*v_avg_j/(64.0*xsection)
scientific_notation="{:.5e}".format(b_binary_i)
print(" Mass of main gas = ", "{:.5e}".format(m_main))
print(" Reduced mass main vs tracer i ", "{:.5e}".format(red_mass_i))
print(" Reduced mass main vs tracer j ", "{:.5e}".format(red_mass_j))
print(" Binary diffusion parameter main gas = ",scientific_notation)
print(" Binary diffusion parameter tracer i = ", "{:.5e}".format(b_binary_i))
print(" Binary diffusion parameter tracer j = ", "{:.5e}".format(b_binary_j))

# Create the array of g(r), and an array of escape velocities for comparisons
g = G*Mp/r**2.0

vescape = np.sqrt(r*g*2.0)

# THERMAL VELOCITIES
v=np.ones(n_array)  #array of 1s
vth_i=sqrt(kb*TK/m_i)*v
vth_j=sqrt(kb*TK/m_j)*v
vth=sqrt(kb*TK/m_main)*v

# INITIALIZE NUMBER DENSITIES
#Initialize the array of number densities using isothermal hydrodynamic equilibrium
ndens=np.ones(n_array)
ndens_i=np.ones(n_array)
ndens_j=np.ones(n_array)
h=kb*TK/(m_main*g)
h_i=kb*TK/(m_i*g)
h_j=kb*TK/(m_j*g)
z_atm=(r-R_s)

# This is the equilibrium initial condition
g_0=G*Mp/rmin**2.0
h_0=kb*TK/(m_main*g_0)
h_0_i=kb*TK/(m_i*g_0)
h_0_j=kb*TK/(m_j*g_0)
ndens_i=mix_i*ndens_initial*np.exp(r/h_i-rmin/h_0_i)
ndens_j=mix_j*ndens_initial*np.exp(r/h_j-rmin/h_0_j)
ndens=ndens_initial*np.exp(r/h-rmin/h_0)


# Alternative initial condition to show that the result
# does not depend on it
for i in range(1,n):
    new_background=1.0e-10
    ndens[i]=new_background
    ndens_i[i]=mix_i*new_background
    ndens_j[i]=mix_j*new_background
    


moles_main=0.0
moles_i=0.0
moles_j=0.0
for i in range(0,n+1):
    moles_main=ndens[i]+ndens[i]*volume[i]
    moles_i=ndens_i[i]+ndens_i[i]*volume[i]
    moles_j=ndens_j[i]+ndens_j[i]*volume[i]
molefraction=moles_main/(moles_main+moles_i+moles_j)
molefraction_i=moles_i/(moles_main+moles_i+moles_j)
molefraction_j=moles_j/(moles_main+moles_i+moles_j)


# Alternative means of tracking iniital evaporation
# composition if trying to use unrealistic initial moles
#molefraction_i=mix_i/(1+mix_i)
#molefraction_j=mix_j/(1+mix_j)
#molefraction=1.0-molefraction_i-molefraction_j

print('')
print("COMPOSITION OF EVAPORATION PRODUCT:")
print("Bulk mole fraction carrier = ", molefraction)
print("Bulk mole fraction i = ",molefraction_i)
print("Bulk mole fraction j = ",molefraction_j)

height=z_atm/1000.0  #for kilometers

# TIME STEP SIZE
# First calculate a diffusivity at initial pressures
print(' ')
Dtemp=b_binary_i/ndens[0] # not really needed now that user inputs dt
print('Diffusion coefficient at base = ',Dtemp,'m^2/s')
print(' ')
print(" dt  = ", dt, "seconds")
print('')

# PRESSURES IN PA AND BAR
P_Pa=(kb*TK)*ndens
P_Pa_i=(kb*TK)*ndens_i
P_Pa_j=(kb*TK)*ndens_j
P_bar=P_Pa*1.0e-5
P_bar_i=P_Pa_i*1.0e-5
P_bar_j=P_Pa_j*1.0e-5

# Check initial conditions using a plot if desired,
# and then after dismissing the plot, prompt user if they wish to continue
#hkm=h/1000.0
#zhratio=height/hkm
#plt.plot(ndens, height)
#plt.xlabel("n")
#plt.ylabel("z (km)")
#plt.show()
#
#go=input('Continue...? (n for no)' )
#if go == 'n':
#    exit()

print(" Pressure at ground level due to evaporation = ", P_Pa[0]," Pa")
print(" ...initial pressure above ground level = ", P_Pa[5]," Pa")
print('')

# DIFFUSION AND GRAVITATIONAL VELOCITIES
# Initial vector of mean free paths for each position, applies to both masses
L=1.0/(xsection*ndens)

# Calculate Knudsen number at each radial position for tracking
kn=L/h

# Vector of terminal velocities for each radial position
wf=g*L/vth
wf_i=g*L/(vth_i)
wf_j=g*L/(vth_j)
print(" TERMINAL VELOCITIES")
print(" Terminal velocity at ground main gas = ","{:.5e}".format(wf[1]))
print(" Terminal velocity at ground tracer i = ", "{:.5e}".format(wf_i[1]))
print(" Terminal velcoity at ground tracer j = ", "{:.5e}".format(wf_j[1]))
print('')

# Initial vectors of diffusivities for each radial position
D_i=b_binary_i/ndens
D_j=b_binary_j/ndens

# DEFINE EDDY DIFFUSION COEFFICIENT, gives an upward flux due to eddy diffusion coefficient
#keddy=1.0e04 #in m^2/s, 10^4 removes negative u at height
D_i=D_i+keddy
D_j=D_j+keddy

# Net velocities are recorded for each radial position
# Create arrays of first derivatives and use Forward differencing
# for first radial position.  These velocities are RELATIVE to the
# mean velocities of the gas.
dndr=np.ones(n_array)
dndr_i=np.ones(n_array)
dndr_j=np.ones(n_array)
dndr[0]=(ndens[1]-ndens[0])/dr
dndr_i[0]=(ndens_i[1]-ndens_i[0])/dr
dndr_j[0]=(ndens_j[1]-ndens_j[0])/dr
dwfdr=np.ones(n_array)
dwfdr_i=np.ones(n_array)
dwfdr_j=np.ones(n_array)
dwfdr[0]=(wf[1]-wf[0])/dr
dwfdr_i[0]=(wf_i[1]-wf_i[0])/dr
dwfdr_j[0]=(wf_j[1]-wf_j[0])/dr
# central difference for every other radial position
i=0
while i < nr-1:
    i=i+1
    dndr[i]=(ndens[i+1]-ndens[i-1])/(2.0*dr)
    dndr_i[i]=(ndens_i[i+1]-ndens[i-1])/(2.0*dr)
    dndr_j[i]=(ndens_j[i+1]-ndens[i-1])/(2.0*dr)
    dwfdr[i]=(wf[i+1]-wf[i-1])/(2.0*dr)
    dwfdr_i[i]=(wf_i[i+1]-wf_i[i-1])/(2.0*dr)
    dwfdr_j[i]=(wf_j[i+1]-wf_j[i-1])/(2.0*dr)
# backward differencing for outer boundary
dndr[n]=(ndens[n]-ndens[n-1])/(dr)
dndr_i[i]=(ndens_i[n]-ndens[n-1])/(dr)
dndr_j[i]=(ndens_j[n]-ndens[n-1])/(dr)
dwfdr[n]=(wf[n]-wf[n-1])/(dr)
dwfdr_i[n]=(wf_i[n]-wf_i[n-1])/(dr)
dwfdr_j[n]=(wf_j[n]-wf_j[n-1])/(dr)
u=-D_i*1.0/ndens*dndr-wf
u_i=-D_i*1.0/ndens_i*dndr_i-wf_i
u_j=-D_j*1.0/ndens_j*dndr_j-wf_j
u_no_grav=-D_i*1/ndens*dndr
    
print(" initial dn/dr[0] = ", dndr[0])
print(" initial dwf/dr_i[0] = ", dwfdr_i[0])
print('')
print(" initial dn/dr[3] = ", dndr[3])
print(" initial dwf/dr_i[3] = ",dwfdr_i[3])
print('')
print(" initial dn/dr[nr] = ", dndr[n])
print(" initial dwf/dr_i[nr] = ", dwfdr_i[n])
print('')

# EQUILIBRIUM VAPOR PRESSURE
PEQ_Pa=PEQ*1.0e5
neq=PEQ_Pa/(kb*TK)
PEQ_Pa_i=PEQ_Pa*mix_i
neq_i=PEQ_Pa_i/(kb*TK)
PEQ_Pa_j=PEQ_Pa*mix_j
neq_j=PEQ_Pa_j/(kb*TK)
print(" ndensity equilibrium = ",neq)
print(" ndensity i equilibrium = ", neq_i)
print(" ndensity j equilibrium = ", neq_j)
print(" initial ndensity = ", ndens_initial)

# Escape parameter lambda at surrace of body at TK
lambda_escape=(G*Mp*m_main/R_s)/(kb*TK)

# Setp matrices for implicit finite difference.
# M is the coefficient matrix, Minv its inverse, x is the new ndens values with
# dimensions n-1 since it omits the two boundary positions (ndens is n+1), and
# y is the forcing function (r.h.s.).
# There are n-1 total positions in each dimension, which is 2 less than for the radius
# vector that has n+1.  Confusing.
M=np.zeros((n-1,n-1))
x=np.zeros(n-1)
y=np.zeros(n-1)
# Vectors of coefficients for M
A=np.zeros(n-1)
B=np.zeros(n-1)
C=np.zeros(n-1)
#
M_i=np.zeros((n-1,n-1))
x_i=np.zeros(n-1)
y_i=np.zeros(n-1)
# Vectors of coefficients for M
A_i=np.zeros(n-1)
B_i=np.zeros(n-1)
C_i=np.zeros(n-1)
#
M_j=np.zeros((n-1,n-1))
x_j=np.zeros(n-1)
y_j=np.zeros(n-1)
# Vectors of coefficients for M
A_j=np.zeros(n-1)
B_j=np.zeros(n-1)
C_j=np.zeros(n-1)

gamma=dt/dr**2
# TIME STEPS BEGIN HERE
#
# INCREMENT NUMBER DENSITY AND THEREFORE PRESSURE, MEAN FREE PATH, TERMINAL VELOCITY, DIFFUSION
# COEFFICIENTS, AND DERIVATIVE WITH RESPECT TO r[ ] FOR EACH TIME STEP
#
print('')
print("TIME STEPS BEGINNING...")
t=0
j=0
#nlimit=50  #NUMBER OF TIME STEPS 100
NREPORT=5
nsurface=np.ones(nlimit+1)
ntop=np.ones(nlimit+1)
timesteps=np.zeros(nlimit+1)
nsurface[0]=ndens[5]
ntop[0]=ndens[n-2]
ndens_old=ndens
n_hydro=n
rc_z=(rc-R_s)/1000.0 # transonic point in km
dndt_inner=0.0
dndt_outer=0.0
dndt_inner_used=0.0
dndt_lower=0.0
dndt_lower_i=0.0
dndt_lower_j=0.0
dndt_outer_used=0.0
dndt_outer_used_i=0.0
dndt_outer_used_j=0.0
seconds_per_year=3.14e7
k=0
success=0


#vel_factor=0.1  # ACCELERATION AT OUTER BOUNDARY 2.2 stable
#esc_factor=0.9  # 0.975 or 0.985 sets a boundary velocity that accerates gas through column
while j < nlimit:
    # Report status from last time step before going on to this time step
    if j%NREPORT == 0:  #For integer values of time step counter report progress
            print('')
            print('')
            print("* New time step ",j,'*')
            print("...Time = ","{:.5e}".format(t),"seconds")
            print("...Time = ","{:.5e}".format(t/seconds_per_year),"years")
            print("...Critical radius for escape = ", "{:.5e}".format(rc)," m")
            print("...Critical radius as height above surface = ","{:.5e}".format(rc_z)," km" )
            print("...Lambda escape at r[o] = ",lambda_escape)
            print("...Sound speed = ", Cs)
            print("...D + Keddy [0] = ","{:.5e}".format(D_i[0]))
            print("...D + Keddy [1] = ","{:.5e}".format(D_i[1]))
            print("...D + Keddy [2] = ","{:.5e}".format(D_i[2]))
            print("...wf[1] = ", "{:.5e}".format(wf[1]))
            print("...wf[200] = ", "{:.5e}".format(wf[200]))
            print("...wf[n] = ", "{:.5e}".format(wf[n]))
            print("...u[0]   = ", "{:.5e}".format(u[0]))
            print("...u[1]   = ", "{:.5e}".format(u[1]))
            print("...u[2]   = ", "{:.5e}".format(u[2]))
            print("...u[50]  = ", "{:.5e}".format(u[50]))
            print("...u[100] = ", "{:.5e}".format(u[100]))
            print("...u[n-1]   = ", "{:.5e}".format(u[n-1]))
            print("...u[n]   = ", "{:.5e}".format(u[n]))
            print("...u[n] no grav = ", "{:.5e}".format(u_no_grav[n]))
            print("...wave position = ", n_hydro)
            print("---------------------------------")
            print("...P_bar[0]   = ",P_bar[0]," bar")
            print("...P_bar[1]   = ",P_bar[1]," bar")
            print("...P_bar[2]   = ",P_bar[2]," bar")
            print("...P_bar[3]   = ",P_bar[3]," bar")
            print("...P_bar[4]   = ",P_bar[4]," bar")
            print("...P_bar[100]  = ",P_bar[100]," bar")
            print("...P_bar[600]  = ",P_bar[600]," bar")
            print("...P_bar[n-4] = ", P_bar[n-4]," bar")
            print("...P_bar[n-3] = ", P_bar[n-3]," bar")
            print("...P_bar[n-2] = ", P_bar[n-2]," bar")
            print("...P_bar[n-1] = ", P_bar[n-1]," bar")
            print("...P_bar[n]   = ", P_bar[n]," bar")
            print("-----------------------------------")
            print("...dndr[n-3] = ", dndr[n-3])
            print("...dndr[n-2] = ", dndr[n-2])
            print("...dndr[n-1] = ", dndr[n-1])
            print("...dndr[n]= ", dndr[n])
            print("-----------------------------------")
            print("...dndt at lower boundary = ", dndt_inner)
            print("...dndt at 2nd REV = ", dndt_lower,"#/s")
            print("...dndt at n-5 REV = ", dndt_outer_used)
            print('')
            print("...dndt 2nd REV tracer i = ", dndt_lower_i,"#/s")
            print("...dndt n-5 REV tracer i = ", dndt_outer_used_i)
            print('')
            print("...dndt 2nd REV tracer j = ", dndt_lower_j,"#/s")
            print("...dndt n-5 REV tracer j = ", dndt_outer_used_j)
            print("-----------------------------------")
            print("Newtons method upper boundary = ",success," steps")
    # Increment time
    j=j+1
    t=t+dt
    timesteps[j]=t
    
    # Save previous number densities so dn/dt can be evaluated with them
    ndens_old=ndens
    ndens_old_i=ndens_i
    ndens_old_j=ndens_j
    
    # UPDATE NUMBER DENSITIES
    #start with assigning coefficients, 1 to n is the range for the radial positions included in the
    # matrix equation
    for i in range (1,n):
        A[i-1]=-gamma*D_i[i]+gamma*D_i[i]*dr/r[i]+gamma*D_i[i+1]/4.0-gamma*D_i[i-1]/4.0+gamma*dr*wf[i]/2.0-gamma*dr*u_h[i]/2.0
        B[i-1]=1.0+2.0*gamma*D_i[i]-(wf[i+1]-wf[i-1])*gamma*dr/2.0+gamma*dr*(u_h[i+1]-u_h[i-1])/2.0
        C[i-1]=-gamma*D_i[i]-gamma*D_i[i]*dr/r[i]-gamma*D_i[i+1]/4.0+gamma*D_i[i-1]/4.0-gamma*wf[i]*dr/2.0+gamma*dr*u_h[i]/2.0
    M[0,0]=B[0]
    M[0,1]=C[0]
    for i in range (1,n-2):
        M[i,i+1]=C[i]
        M[i,i]=B[i]
        M[i,i-1]=A[i]
    M[n-2,n-2]=B[n-2]
    M[n-2,n-3]=A[n-2]
    #print("M=",M)
    
    # Setup forcing functon y
    for i in range (0,n-2):
        y[i]=ndens_old[i+1]
    
    # Modify first and last terms for y to account for boundary conditions
    y[0]=y[0]-ndens_old[0]*(-gamma*D_i[1]+gamma*D_i[1]*dr/r[1]+gamma*D_i[2]/4.0-gamma*D_i[0]/4.0+gamma*dr*wf[1]/2.0-gamma*dr*u_h[1]/2.0)
    y[n-2]=y[n-2]-ndens_old[n]*(-gamma*D_i[n-1]-gamma*D_i[n-1]*dr/r[n-1]-gamma*D_i[n]/4.0+gamma*D_i[n-2]/4.0-gamma*wf[n-1]*dr/2.0+gamma*dr*u_h[n-1]/2.0 )
    
    # Solve for x, the number densities between 1 and n-1, with the boundary values at i=0 and i=n
    x=np.dot(np.linalg.inv(M),y)
    
    # Store values in new estimate of number densities ndens
    for i in range (0,n-2):
        ndens[i+1]=x[i]
        
        
    # Inner boundary condition due to evaporation, seems to go too fast, so option to keep constant at
    #
    #equilibrium number density corresponding to equilibrium vapor pressure
    #ndens[0]=ndens_old[0]+dt*area[0]*((neq-ndens_old[0])*sqrt(kb*TK)/(sqrt(2.0*pi*red_mass)))/volume[0]
    dndt_inner=ndens_old[0]*u[0]*area[0]
    dndt_lower=ndens_old[1]*u[1]*area[1]
    ndens[0]=neq  #Simplest lower boundary condition that others converge to rapidlly
       
    dndt_outer_used=ndens_old[n-5]*u[n-5]*area[n-5]
    
    # OUTER BOUNDARY CONDITION due to hydrodynamic escape. Trying various options for numerical stability
    dndt_outer=u[n]*ndens[n]*area[n]
    ndens[n]=vel_factor*ndens_old[n]   #This one is reliable, continuity by slow leak
    
    # Newtons method for flux boundary condition
    precision=1.0e+01
    ntest=ndens_old[n]
    fprime=-u[n-1]*4.0*pi*r[n]**2.0
    success=0
    for k in range(0,100):
        ftest=u[n-1]*area[n-1]*ndens_old[n-1]-vel_factor*u[n-1]*area[n]*ntest
        if abs(ftest) < precision:
            ndens[n]=ntest
        else:
            ntest=ntest-ftest/fprime
            success = success+1
        k = k+1
    ndens[n]=ntest
    if ndens[n] < 0.0:
        ndens[n]=1.0e-20
    
    n_hydro=n
   
    # Update pressures
    P_Pa=(kb*TK)*ndens
    P_bar=P_Pa*1.0e-5
    
    
    # UPDATE TOTOAL NUMBER DENSITIES FOR i
    #start with assigning coefficients, 1 to n is the range for the radial positions included in the
    # matrix equation
    for i in range (1,n):
        A_i[i-1]=-gamma*D_i[i]+gamma*D_i[i]*dr/r[i]+gamma*D_i[i+1]/4.0-gamma*D_i[i-1]/4.0+gamma*dr*wf_i[i]/2.0-gamma*dr*u_h_i[i]/2.0
        B_i[i-1]=1.0+2.0*gamma*D_i[i]-(wf_i[i+1]-wf_i[i-1])*gamma*dr/2.0+gamma*dr*(u_h_i[i+1]-u_h_i[i-1])/2.0
        C_i[i-1]=-gamma*D_i[i]-gamma*D_i[i]*dr/r[i]-gamma*D_i[i+1]/4.0+gamma*D_i[i-1]/4.0-gamma*wf_i[i]*dr/2.0+gamma*dr*u_h_i[i]/2.0
    M_i[0,0]=B_i[0]
    M_i[0,1]=C_i[0]
    for i in range (1,n-2):
        M_i[i,i+1]=C_i[i]
        M_i[i,i]=B_i[i]
        M_i[i,i-1]=A_i[i]
    M_i[n-2,n-2]=B_i[n-2]
    M_i[n-2,n-3]=A_i[n-2]
    #print("M=",M)
    
    # Setup forcing functon y
    for i in range (0,n-2):
        y_i[i]=ndens_old_i[i+1]
    
    # Modify first and last terms for y to account for boundary conditions
    y_i[0]=y_i[0]-ndens_old_i[0]*(-gamma*D_i[1]+gamma*D_i[1]*dr/r[1]+gamma*D_i[2]/4.0-gamma*D_i[0]/4.0+gamma*dr*wf_i[1]/2.0-gamma*dr*u_h_i[1]/2.0)
    y_i[n-2]=y_i[n-2]-ndens_old_i[n]*(-gamma*D_i[n-1]-gamma*D_i[n-1]*dr/r[n-1]-gamma*D_i[n]/4.0+gamma*D_i[n-2]/4.0-gamma*wf_i[n-1]*dr/2.0+gamma*dr*u_h_i[n-1]/2.0)
    
    # Solve for x, the number densities between 1 and n-1, with the boundary values at i=0 and i=n
    x_i=np.dot(np.linalg.inv(M_i),y_i)
    
    # Store values in new estimate of number densities ndens
    for i in range (0,n-2):
        ndens_i[i+1]=x_i[i]
        
    # Inner boundary condition due to evaporation
    dndt_lower_i=ndens_old_i[1]*u[1]*area[1]
    ndens_i[0]=neq_i
       
    dndt_outer_used_i=ndens_old_i[n-5]*u[n-5]*area[n-5]
    
    # OUTER BOUNDARY CONDITION due to hydrodynamic escape.
    #ndens_i[n]=vel_factor*ndens_old_i[n]
    
    # Newtons method for flux boundary condition
    precision=1.0e+01
    ntest=ndens_old_i[n]
    fprime=-u_i[n-1]*4.0*pi*r[n]**2.0
    success=0
    for k in range(0,100):
        ftest=u_i[n-1]*area[n-1]*ndens_old_i[n-1]-vel_factor*u_i[n-1]*area[n]*ntest
        if abs(ftest) < precision:
            ndens_i[n]=ntest
        else:
            ntest=ntest-ftest/fprime
            success = success+1
        k = k+1
    ndens_i[n]=ntest
    if ndens_i[n] < 0.0:
        ndens_i[n]=1.0e-20
 
    # Update pressures
    P_Pa_i=(kb*TK)*ndens_i
    P_bar_i=P_Pa_i*1.0e-5
    
    
    # UPDATE TOTOAL NUMBER DENSITIES FOR j
    #start with assigning coefficients, 1 to n is the range for the radial positions included in the
    # matrix equation
    for i in range (1,n):
        A_j[i-1]=-gamma*D_j[i]+gamma*D_j[i]*dr/r[i]+gamma*D_j[i+1]/4.0-gamma*D_j[i-1]/4.0+gamma*dr*wf_j[i]/2.0-gamma*dr*u_h_j[i]/2.0
        B_j[i-1]=1.0+2.0*gamma*D_j[i]-(wf_j[i+1]-wf_j[i-1])*gamma*dr/2.0+gamma*dr*(u_h_j[i+1]-u_h_j[i-1])/2.0
        C_j[i-1]=-gamma*D_j[i]-gamma*D_j[i]*dr/r[i]-gamma*D_j[i+1]/4.0+gamma*D_j[i-1]/4.0-gamma*wf_j[i]*dr/2.0+gamma*dr*u_h_j[i]/2.0
    M_j[0,0]=B_j[0]
    M_j[0,1]=C_j[0]
    for i in range (1,n-2):
        M_j[i,i+1]=C_j[i]
        M_j[i,i]=B_j[i]
        M_j[i,i-1]=A_j[i]
    M_j[n-2,n-2]=B_j[n-2]
    M_j[n-2,n-3]=A_j[n-2]
    #print("M=",M)
    
    # Setup forcing functon y
    for i in range (0,n-2):
        y_j[i]=ndens_old_j[i+1]
    
    # Modify first and last terms for y to account for boundary conditions
    y_j[0]=y_j[0]-ndens_old_j[0]*(-gamma*D_j[1]+gamma*D_j[1]*dr/r[1]+gamma*D_j[2]/4.0-gamma*D_j[0]/4.0+gamma*dr*wf_j[1]/2.0-gamma*dr*u_h_j[1]/2.0)
    y_j[n-2]=y_j[n-2]-ndens_old_j[n]*(-gamma*D_j[n-1]-gamma*D_j[n-1]*dr/r[n-1]-gamma*D_j[n]/4.0+gamma*D_j[n-2]/4.0-gamma*wf_j[n-1]*dr/2.0+gamma*dr*u_h_j[n-1]/2.0)
    
    # Solve for x, the number densities between 1 and n-1, with the boundary values at i=0 and i=n
    x_j=np.dot(np.linalg.inv(M_j),y_j)
    
    # Store values in new estimate of number densities ndens
    for i in range (0,n-2):
        ndens_j[i+1]=x_j[i]
        
    # Inner boundary condition due to evaporation
    dndt_lower_j=ndens_old_j[1]*u[1]*area[1]  #this is molecules/s not density
    ndens_j[0]=neq_j
       
    dndt_outer_used_j=ndens_old_j[n-5]*u[n-5]*area[n-5]
    
    # OUTER BOUNDARY CONDITION due to hydrodynamic escape.
    #ndens_j[n]=vel_factor*ndens_old_j[n]
    
    # Newtons method for flux boundary condition
    precision=1.0e+01
    ntest=ndens_old_j[n]
    fprime=-u_j[n-1]*4.0*pi*r[n]**2.0
    success=0
    for k in range(0,100):
        ftest=u_j[n-1]*area[n-1]*ndens_old_j[n-1]-vel_factor*u_j[n-1]*area[n]*ntest
        if abs(ftest) < precision:
            ndens_j[n]=ntest
        else:
            ntest=ntest-ftest/fprime
            success = success+1
        k = k+1
    ndens_j[n]=ntest
    if ndens_j[n] < 0.0:
        ndens_j[n]=1.0e-20

    # Update pressures
    P_Pa_j=(kb*TK)*ndens_j
    P_bar_j=P_Pa_j*1.0e-5
    
    # Update L, diffusivities, terminal velocities, and knudsen numbers
    L=1.0/(xsection*ndens)
    D_i=(b_binary_i/ndens)+keddy
    D_j=(b_binary_j/ndens)+keddy
    wf=g*L/(vth)
    wf_i=g*L/(vth_i)
    wf_j=g*L/(vth_j)
    kn=L/h
    
    
    # UPDATE DERIVATIVES -
    dndr[0]=(ndens[1]-ndens[0])/dr
    dndr_i[0]=(ndens_i[1]-ndens_i[0])/dr
    dndr_j[0]=(ndens_j[1]-ndens_j[0])/dr
    dwfdr[0]=(wf[1]-wf[0])/dr
    dwfdr_i[0]=(wf_i[1]-wf_i[0])/dr
    dwfdr_j[0]=(wf_j[1]-wf_j[0])/dr
    # central difference for every other radial position
    i=0
    while i < nr-1:
        i=i+1
        dndr[i]=(ndens[i+1]-ndens[i-1])/(2.0*dr)
        dndr_i[i]=(ndens_i[i+1]-ndens[i-1])/(2.0*dr)
        dndr_j[i]=(ndens_j[i+1]-ndens[i-1])/(2.0*dr)
        dwfdr[i]=(wf[i+1]-wf[i-1])/(2.0*dr)
        dwfdr_i[i]=(wf_i[i+1]-wf_i[i-1])/(2.0*dr)
        dwfdr_j[i]=(wf_j[i+1]-wf_j[i-1])/(2.0*dr)
    # backward differencing for outer boundary
    dndr[n]=(ndens[n]-ndens[n-1])/(dr)
    dndr_i[i]=(ndens_i[n]-ndens[n-1])/(dr)
    dndr_j[i]=(ndens_j[n]-ndens[n-1])/(dr)
    dwfdr[n]=(wf[n]-wf[n-1])/(dr)
    dwfdr_i[n]=(wf_i[n]-wf_i[n-1])/(dr)
    dwfdr_j[n]=(wf_j[n]-wf_j[n-1])/(dr)
    
    #Record near surface number density
    nsurface[j]=ndens[5]
    ntop[j]=ndens[n-2]
    
    # Update velocities
    u=u_h-D_i*1.0/ndens*dndr-wf
    u_i=u_h_i-D_i*1.0/ndens_i*dndr_i-wf_i
    u_j=u_h_j-D_j*1.0/ndens_j*dndr_j-wf_j
    u_no_grav=u_h-D_i*1/ndens*dndr
    vel_escape_ratio=u/vescape
    vel_sound_ratio=u/Cs
    Cs_escape_ratio=Cs/vescape
    
# END TIME STEP

print("Time steps completed, making plots:")

# Counting atoms, not moles actually
moles_main=0.0
moles_i=0.0
moles_j=0.0
for i in range(0,n+1):
    moles_main=ndens[i]+ndens[i]*volume[i]
    moles_i=ndens_i[i]+ndens_i[i]*volume[i]
    moles_j=ndens_j[i]+ndens_j[i]*volume[i]
molefraction=moles_main/(moles_main+moles_i+moles_j)
molefraction_i=moles_i/(moles_main+moles_i+moles_j)
molefraction_j=moles_j/(moles_main+moles_i+moles_j)
print('')
print("FINAL ATMOSPHERE COMPOSITION:")
print("Bulk mole fraction main = ", molefraction)
print("Bulk mole fraction i = ",molefraction_i)
print("Bulk mole fraction j = ",molefraction_j)
print('')

# Energy calculations
totalatoms=moles_main+moles_i+moles_j
totalmass=totalatoms*m_main
Energy=1.0/(1.666-1.0)*kb*(totalatoms)*TK
index_rc=0
for i in range (0,n):
    if r[i] < rc:
        index_rc=index_rc+1
ndot=ndens[index_rc]*Cs*4.0*pi*r[index_rc]**2
energy_dot=(G*Mp*ndot/rc)*m_main +(kb*TK/ndens[index_rc])*ndot
thermal_dot=(kb*TK/ndens[index_rc])*ndot
ndot_in=ndens[1]*u[1]*area[1]
energy_in_dot=(kb*TK/ndens[0])*ndot_in
CV=Energy/TK
print('At critical point r = ',r[index_rc])
print('Crtical radius = ',rc,' Should match above')
print('Pressure at critical point = ',P_bar[index_rc],' bar')
print('Total atoms comprising atmosphere = ',totalatoms)
print('Total mass of atm in kg = ',totalmass)
print('Total thermal Energy in Atmosphere = ',Energy," Joules")
print('Total heat capacity of atmosphere = ', "{:.5e}".format(CV),' J/K')
print('atoms per second at rc = ',ndot)
print('Joules per second due to atmospheric loss = ', "{:.5e}".format(energy_dot))
print('Joules per second thermal energy lost at rc = ',"{:.5e}".format(thermal_dot))
print('')
print('')

# MAKE PLOTS
logP=np.log10(mix*P_bar)
logP_i=np.log10(P_bar_i)
logP_j=np.log10(P_bar_j)
ratio=((P_bar_j/P_bar_i)/(mix_j/(mix_i))-1.0) #delta but not in per mil
ratio_i=(ndens_i/(ndens))/(mix_i)-1.0
mixing_ratio_i=(ndens_i)/(ndens+ndens_i+ndens_j)

a_file = open('isotope_ratio.txt', 'w')
for i in range(0,n):
    a_file.write("%13.8f\n" % ratio[i])
a_file.close()

a_file = open('height_km.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5f\n" % height[i])
a_file.close()

a_file = open('velocity.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5f\n" % u[i])
a_file.close()

a_file = open('mixing_ratio.txt', 'w')
for i in range(0,n):
    a_file.write("%11.6f\n" % mixing_ratio_i[i])
a_file.close()

a_file = open('tracer_relative_ratio.txt', 'w')
for i in range(0,n):
    a_file.write("%13.8f\n" % ratio_i[i])
a_file.close()

a_file = open('pressure_main.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % P_bar[i])
a_file.close()

a_file = open('pressure_i.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % P_bar_i[i])
a_file.close()

a_file = open('scaleheight_main_gas_meters.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % h[i])
a_file.close()

a_file = open('wf.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % wf[i])
a_file.close()

a_file = open('u_h.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % u_h[i])
a_file.close()

a_file = open('u_h_i.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % u_h_i[i])
a_file.close()

a_file = open('n_density.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % ndens[i])
a_file.close()

#a_file = open('Cs.txt', 'w')
#for i in range(0,n):
#    a_file.write("%10.5e\n" % Cs[i])
#a_file.close()

a_file = open('vescape.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % vescape[i])
a_file.close()

a_file = open('volumes.txt', 'w')
for i in range(0,n):
    a_file.write("%10.5e\n" % volume[i])
a_file.close()

a_file = open('knudsen.txt', 'w')
for i in range(0,n):
    a_file.write("%12.7e\n" % kn[i])
a_file.close()

nmaxplot=n-20

# Plot pressure of major gas species and first tracer
plt.figure(1)
plt.plot(logP[0:nmaxplot],height[0:nmaxplot],label=m_main_string)
plt.plot(logP_i[0:nmaxplot],height[0:nmaxplot],label=mi_string)
plt.title('PRESSURE')
plt.xlabel('log(Pressure) (bar))')
plt.ylabel(' Z (km)')
plt.ylim(height[0],heightrc)
# Add a Legend
plt.legend()
# Show the plot
#plt.show()

# Plot mixing mole fraction i ratio
plt.figure(2)
plt.plot(mixing_ratio_i[0:nmaxplot],height[0:nmaxplot],label=mi_string)
plt.ylim(height[0],heightrc)
plt.title('Mixing ratio of' + '  ' + mi_string )
plt.xlabel('Mixing ratio ' + ' ' + mi_string )
plt.ylabel(' Z (km)')
plt.legend()
# Show the plot
#plt.show()

# Plot i/main ratio
plt.figure(3)
plt.plot(ratio_i[0:nmaxplot],height[0:nmaxplot],label=mi_string)
plt.ylim(height[0],heightrc)
plt.title('Ratio (n_i/n)/initial -1' + '  ' + mi_string + '/' + m_main_string  )
plt.xlabel('xi/x_initial -1' )
plt.ylabel(' Z (km)')
plt.legend()
# Show the plot
#plt.show()

# Isotope ratio plot, j/i
plt.figure(4)
plt.plot(ratio[0:nmaxplot], height[0:nmaxplot])
plt.ylim(height[0],heightrc)
plt.title('Ratio (n_j/n_i)/initial -1' + '  ' + mj_string + '/' + mi_string  )
plt.xlabel('Ratio/Ratio_i -1')
plt.ylabel('Z (km)')
#plt.xlim(0.0,maxht-2.0*dr)
# Add a grid
#plt.grid(alpha=.4,linestyle='--')
#Add a Legend
#plt.legend()
#plt.xlim(0.0,1.5)
#plt.ylim(0.0,height[n-10])
## show plot
#plt.show()

# Velocities plot
Cs_array=np.zeros(n+1)
for i in range(0,n):
    Cs_array[i]=Cs
plt.figure(5)
plt.plot(u[0:nmaxplot], height[0:nmaxplot])
plt.plot(vescape[0:nmaxplot],height[0:nmaxplot],linestyle='--',label='Escape vel')
plt.plot(Cs_array[0:nmaxplot],height[0:nmaxplot],linestyle='-.',label='Cs')
plt.plot(u_no_grav[0:nmaxplot],height[0:nmaxplot],linestyle=':',label='no gravity')
plt.ylim(height[0],heightrc)
plt.xlabel("velocity (m/s)")
plt.ylabel("height (km)")
plt.legend()
#plt.show()

plt.figure(6)
plt.plot(vel_escape_ratio[0:nmaxplot], height[0:nmaxplot],label='u/u_escape')
plt.plot(vel_sound_ratio[0:nmaxplot], height[0:nmaxplot],label='u/Cs')
plt.plot(Cs_escape_ratio[0:nmaxplot], height[0:nmaxplot],linestyle='--',label='Cs/u_escape')
plt.ylim(height[0],heightrc)
plt.xlabel("u/u_i")
plt.ylabel("height (km)")
#plt.xlim(-3.0,3.0)
plt.legend()

plt.figure(7)
plt.plot(kn[0:nmaxplot], height[0:nmaxplot])
plt.ylim(height[0],heightrc)
plt.xlabel("Knudsen number")
plt.ylabel("height (km)")
#plt.xlim(-3.0,3.0)

# log velocity
plt.figure(9)
logu=np.log10(u)
logui=np.log10(u_i)
plt.plot(logu[0:nmaxplot], height[0:nmaxplot],label='Main species')
plt.plot(logui[0:nmaxplot],height[0:nmaxplot],linestyle=':',label='Minor species')
plt.ylim(height[0],heightrc)
plt.xlabel("log velocity (m/s)")
plt.ylabel("height (km)")
plt.legend()
#plt.show()

r_ratio_R_s=(r/R_s)
uratio=(u/Cs)
plt.figure(10)
plt.plot(r_ratio_R_s[0:nmaxplot], uratio[0:nmaxplot], label = 'Main species')
plt.axhline(y=1.00,linestyle='--')
plt.axvline(x=rc_ratio,linestyle='--',label='rc')
plt.xlabel("r/r_planet")
plt.ylabel("Velocity/Cs")
plt.legend()
plt.yscale('log')
plt.xscale('log')
#plt.ylim(1.0e-10,2.0)
plt.title('Lambda = ' + lambda_string)

plt.show()  # Including only one plt.show command makes all plots in their own windows

#plt.plot(ratio, height)
#plt.title('Ratio m_j/m_i')
#plt.xlabel('Ratio')
#plt.ylabel('Z (km)')
#plt.xlim(0.0,maxht-2.0*dr)
# Add a grid
#plt.grid(alpha=.4,linestyle='--')
# Add a Legend
#plt.legend()
# Show the plot
#plt.xlim(0.0,1.5)
#plt.ylim(0.0,height[n-10])
#plt.show()

# Create the plot to check value
#lognsurface=np.log10(nsurface)
#logntop=np.log10(ntop)
#plt.plot(timesteps,lognsurface,label='near surface')
#plt.plot(timesteps,logntop,label='top')
#plt.title('n near surface')
#plt.xlabel('Time (s)')
#plt.ylabel('log(Number density)')
#plt.xlim(0.0,maxht-2.0*dr)
# Add a grid
#plt.grid(alpha=.4,linestyle='--')
# Add a Legend
#plt.legend()
# Show the plot
#plt.xlim(-10.0,10.0)
#plt.show()


print('Done...')
exit()



